package com.datareport.compute;

import com.datareport.common.exception.BusinessException;
import com.datareport.config.FlinkComputeConfig;
import com.datareport.domain.entity.StatRule;
import com.datareport.repository.StatRuleRepository;
import com.datareport.service.FlinkComputeService;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.JobStatus;
import org.apache.flink.client.program.ClusterClient;
import org.apache.flink.client.program.rest.RestClusterClient;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.RestOptions;
import org.apache.flink.runtime.messages.Acknowledge;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * Flink作业管理器
 * 管理Flink作业的生命周期，包括启动、停止、监控等
 */
@Slf4j
@Component
public class FlinkJobManager {

    @Autowired
    private FlinkComputeService computeService;

    @Autowired
    private FlinkComputeConfig computeConfig;

    @Autowired
    private StatRuleRepository ruleRepository;

    // 作业状态缓存
    private final Map<String, JobStatusInfo> jobStatusCache = new ConcurrentHashMap<>();

    // Flink集群客户端
    private ClusterClient<String> clusterClient;

    /**
     * 初始化Flink集群客户端
     */
    @PostConstruct
    public void init() {
        try {
            Configuration flinkConfig = new Configuration();
            flinkConfig.setString(RestOptions.ADDRESS, "localhost");
            flinkConfig.setInteger(RestOptions.PORT, 8081);
            
            // 创建集群客户端
            clusterClient = new RestClusterClient<>(flinkConfig, "standalone");
            
            log.info("Flink集群客户端初始化成功");
        } catch (Exception e) {
            log.error("Flink集群客户端初始化失败", e);
            throw new RuntimeException("Flink集群客户端初始化失败", e);
        }
    }

    /**
     * 销毁资源
     */
    @PreDestroy
    public void destroy() {
        if (clusterClient != null) {
            try {
                clusterClient.close();
                log.info("Flink集群客户端已关闭");
            } catch (Exception e) {
                log.error("关闭Flink集群客户端失败", e);
            }
        }
    }

    /**
     * 启动规则计算作业
     * @param ruleId 规则ID
     * @return 作业ID
     */
    public String startJob(String ruleId) {
        try {
            // 检查规则是否存在且启用
            StatRule rule = ruleRepository.selectById(ruleId);
            if (rule == null) {
                throw new BusinessException("规则不存在: " + ruleId);
            }
            
            if (rule.getStatus() != 1) {
                throw new BusinessException("规则未启用: " + ruleId);
            }

            // 检查作业是否已运行
            if (isJobRunning(ruleId)) {
                throw new BusinessException("作业已在运行中: " + ruleId);
            }

            log.info("开始启动Flink作业: ruleId={}", ruleId);

            // 异步启动作业
            CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {
                try {
                    computeService.startComputeJob(ruleId);
                } catch (Exception e) {
                    log.error("Flink作业启动失败: ruleId={}", ruleId, e);
                    updateJobStatus(ruleId, JobStatus.FAILED, e.getMessage());
                }
            });

            // 等待作业启动完成
            try {
                future.get(30, TimeUnit.SECONDS);
                updateJobStatus(ruleId, JobStatus.RUNNING, "作业启动成功");
                log.info("Flink作业启动成功: ruleId={}", ruleId);
                return ruleId;
            } catch (Exception e) {
                log.error("Flink作业启动超时: ruleId={}", ruleId, e);
                updateJobStatus(ruleId, JobStatus.FAILED, "作业启动超时");
                throw new BusinessException("作业启动超时: " + ruleId);
            }

        } catch (BusinessException e) {
            throw e;
        } catch (Exception e) {
            log.error("启动Flink作业失败: ruleId={}", ruleId, e);
            throw new BusinessException("启动作业失败: " + e.getMessage());
        }
    }

    /**
     * 停止规则计算作业
     * @param ruleId 规则ID
     */
    public void stopJob(String ruleId) {
        try {
            JobStatusInfo statusInfo = jobStatusCache.get(ruleId);
            if (statusInfo == null || !statusInfo.isRunning()) {
                log.warn("作业未运行，无需停止: ruleId={}", ruleId);
                return;
            }

            String jobId = statusInfo.getJobId();
            if (jobId != null) {
                // 通过Flink API停止作业
                CompletableFuture<Acknowledge> cancelFuture = clusterClient.cancel(JobID.fromHexString(jobId));
                cancelFuture.get(30, TimeUnit.SECONDS);
                
                updateJobStatus(ruleId, JobStatus.CANCELED, "作业已停止");
                log.info("Flink作业停止成功: ruleId={}, jobId={}", ruleId, jobId);
            } else {
                updateJobStatus(ruleId, JobStatus.CANCELED, "作业已停止");
                log.info("Flink作业状态更新为已停止: ruleId={}", ruleId);
            }

        } catch (Exception e) {
            log.error("停止Flink作业失败: ruleId={}", ruleId, e);
            throw new BusinessException("停止作业失败: " + e.getMessage());
        }
    }

    /**
     * 重启规则计算作业
     * @param ruleId 规则ID
     * @return 新的作业ID
     */
    public String restartJob(String ruleId) {
        try {
            log.info("开始重启Flink作业: ruleId={}", ruleId);
            
            // 先停止现有作业
            stopJob(ruleId);
            
            // 等待一段时间确保作业完全停止
            Thread.sleep(5000);
            
            // 重新启动作业
            return startJob(ruleId);
            
        } catch (BusinessException e) {
            throw e;
        } catch (Exception e) {
            log.error("重启Flink作业失败: ruleId={}", ruleId, e);
            throw new BusinessException("重启作业失败: " + e.getMessage());
        }
    }

    /**
     * 获取作业状态
     * @param ruleId 规则ID
     * @return 作业状态信息
     */
    public JobStatusInfo getJobStatus(String ruleId) {
        JobStatusInfo statusInfo = jobStatusCache.get(ruleId);
        if (statusInfo == null) {
            return new JobStatusInfo(ruleId, JobStatus.FAILED, "作业状态未知");
        }

        // 如果作业正在运行，尝试从Flink集群获取最新状态
        if (statusInfo.isRunning() && statusInfo.getJobId() != null) {
            try {
                JobID jobId = JobID.fromHexString(statusInfo.getJobId());
                CompletableFuture<JobStatus> jobStatusFuture = clusterClient.getJobStatus(jobId);
                JobStatus currentStatus = jobStatusFuture.get(10, TimeUnit.SECONDS);
                
                if (currentStatus != statusInfo.getStatus()) {
                    updateJobStatus(ruleId, currentStatus, "状态更新");
                    statusInfo = jobStatusCache.get(ruleId);
                }
            } catch (Exception e) {
                log.warn("从Flink集群获取作业状态失败: ruleId={}", ruleId, e);
            }
        }

        return statusInfo;
    }

    /**
     * 检查作业是否正在运行
     * @param ruleId 规则ID
     * @return 是否正在运行
     */
    public boolean isJobRunning(String ruleId) {
        JobStatusInfo statusInfo = jobStatusCache.get(ruleId);
        return statusInfo != null && statusInfo.isRunning();
    }

    /**
     * 获取所有作业状态
     * @return 所有作业状态
     */
    public Map<String, JobStatusInfo> getAllJobStatuses() {
        return new ConcurrentHashMap<>(jobStatusCache);
    }

    /**
     * 更新作业状态
     * @param ruleId 规则ID
     * @param status 作业状态
     * @param message 状态消息
     */
    private void updateJobStatus(String ruleId, JobStatus status, String message) {
        JobStatusInfo statusInfo = new JobStatusInfo(ruleId, status, message);
        jobStatusCache.put(ruleId, statusInfo);
        log.info("更新作业状态: ruleId={}, status={}, message={}", ruleId, status, message);
    }

    /**
     * 更新作业状态（包含作业ID）
     * @param ruleId 规则ID
     * @param jobId 作业ID
     * @param status 作业状态
     * @param message 状态消息
     */
    public void updateJobStatus(String ruleId, String jobId, JobStatus status, String message) {
        JobStatusInfo statusInfo = new JobStatusInfo(ruleId, jobId, status, message);
        jobStatusCache.put(ruleId, statusInfo);
        log.info("更新作业状态: ruleId={}, jobId={}, status={}, message={}", ruleId, jobId, status, message);
    }

    /**
     * 清理作业状态
     * @param ruleId 规则ID
     */
    public void clearJobStatus(String ruleId) {
        jobStatusCache.remove(ruleId);
        log.info("清理作业状态: ruleId={}", ruleId);
    }

    /**
     * 获取Flink集群信息
     * @return 集群信息
     */
    public ClusterInfo getClusterInfo() {
        try {
            ClusterInfo info = new ClusterInfo();
            info.setConnected(clusterClient != null);
            
            if (clusterClient != null) {
                // 获取集群概览信息
                String webInterfaceURL = clusterClient.getWebInterfaceURL();
                info.setWebInterfaceUrl(webInterfaceURL);
                info.setClusterId(clusterClient.getClusterId());
            }
            
            return info;
        } catch (Exception e) {
            log.error("获取Flink集群信息失败", e);
            ClusterInfo info = new ClusterInfo();
            info.setConnected(false);
            info.setErrorMessage(e.getMessage());
            return info;
        }
    }

    /**
     * 作业状态信息类
     */
    public static class JobStatusInfo {
        private final String ruleId;
        private String jobId;
        private JobStatus status;
        private String message;
        private long updateTime;

        public JobStatusInfo(String ruleId, JobStatus status, String message) {
            this.ruleId = ruleId;
            this.status = status;
            this.message = message;
            this.updateTime = System.currentTimeMillis();
        }

        public JobStatusInfo(String ruleId, String jobId, JobStatus status, String message) {
            this.ruleId = ruleId;
            this.jobId = jobId;
            this.status = status;
            this.message = message;
            this.updateTime = System.currentTimeMillis();
        }

        public boolean isRunning() {
            return status == JobStatus.RUNNING || status == JobStatus.INITIALIZING || status == JobStatus.RECONCILING;
        }

        // Getters
        public String getRuleId() { return ruleId; }
        public String getJobId() { return jobId; }
        public JobStatus getStatus() { return status; }
        public String getMessage() { return message; }
        public long getUpdateTime() { return updateTime; }

        // Setters
        public void setJobId(String jobId) { this.jobId = jobId; }
        public void setStatus(JobStatus status) { this.status = status; }
        public void setMessage(String message) { this.message = message; }
        public void setUpdateTime(long updateTime) { this.updateTime = updateTime; }
    }

    /**
     * 集群信息类
     */
    public static class ClusterInfo {
        private boolean connected;
        private String clusterId;
        private String webInterfaceUrl;
        private String errorMessage;

        // Getters and setters
        public boolean isConnected() { return connected; }
        public void setConnected(boolean connected) { this.connected = connected; }
        public String getClusterId() { return clusterId; }
        public void setClusterId(String clusterId) { this.clusterId = clusterId; }
        public String getWebInterfaceUrl() { return webInterfaceUrl; }
        public void setWebInterfaceUrl(String webInterfaceUrl) { this.webInterfaceUrl = webInterfaceUrl; }
        public String getErrorMessage() { return errorMessage; }
        public void setErrorMessage(String errorMessage) { this.errorMessage = errorMessage; }
    }
}